# -*- coding: utf-8 -*-
# file: train_keras_resnet.py
# author: JinTian
# time: 15/07/2017 5:11 PM
# Copyright 2017 JinTian. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ------------------------------------------------------------------------
"""
classify Weather data using ResNet50
"""
from keras.applications.resnet50 import ResNet50
from keras.preprocessing import image
from keras.applications.resnet50 import preprocess_input, decode_predictions
import numpy as np
import argparse
from keras.models import Model
from keras.layers import Dense, GlobalAveragePooling2D
from keras import backend as K
from keras.preprocessing.image import ImageDataGenerator


def build_model():
    model = ResNet50(weights='imagenet', include_top=False)
    x = model.output
    x = GlobalAveragePooling2D()(x)
    x = Dense(1024, activation='relu')(x)
    predictions = Dense(4, activation='softmax')(x)
    model = Model(inputs=model.input,
                  outputs=predictions)
    return model


def train(args):
    model = build_model()

    for layer in model.layers:
        layer.trainable = False
    model.compile(optimizer='rmsprop',
                  loss='categorical_crossentropy')

    train_datagen = ImageDataGenerator(
        rescale=1/255,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True,
    )
    train_generator = train_datagen.flow_from_directory(
        args.data_dir,
        target_size=(250, 250),
        batch_size=32,
        class_mode='categorical'
    )
    model.fit_generator(
        train_generator,
        steps_per_epoch=2000,
        epochs=50,
    )


def predict():
    print('this is predict')
    pass


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', '-d', required=True, help='directory contains n classes folders inside'
                                                                'are images.')
    parser.add_argument('--is_train', '-t', type=bool, default=True, help='true to train, false to predict.')

    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    if args.is_train:
        train(args)
    else:
        predict()


if __name__ == '__main__':
    train()
